package com.xp.ai.embedding;


import com.xp.ai.util.ApiKey;
import dev.langchain4j.community.store.embedding.clickhouse.ClickHouseEmbeddingStore;
import dev.langchain4j.community.store.embedding.clickhouse.ClickHouseSettings;
import dev.langchain4j.data.embedding.Embedding;
import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.model.embedding.EmbeddingModel;
import dev.langchain4j.model.openai.OpenAiEmbeddingModel;
import dev.langchain4j.model.output.Response;
import dev.langchain4j.store.embedding.EmbeddingMatch;
import dev.langchain4j.store.embedding.EmbeddingSearchRequest;
import dev.langchain4j.store.embedding.EmbeddingSearchResult;
import dev.langchain4j.store.embedding.EmbeddingStore;
import dev.langchain4j.store.embedding.inmemory.InMemoryEmbeddingStore;
import dev.langchain4j.store.embedding.milvus.MilvusEmbeddingStore;
import dev.langchain4j.store.embedding.pgvector.PgVectorEmbeddingStore;

import java.util.List;

/**
 * 向量相关demo
 * 另一种大模型
 */
public class VectorDemo {


    public static void main(String[] args) {

        //初始化向量模型 和初始化对话模型一样
        //在这里我就明白了，为什么还要单独整一个Model 出来，因为这个不是数据语言大模型了，这是输入嵌入的向量大模型
        //这个向量模型，不是用来做对话的，而是用来做向量搜索的
        EmbeddingModel embeddingModel = OpenAiEmbeddingModel.builder()
                .baseUrl(ApiKey.GJ_BASE_URL)
                .apiKey(ApiKey.GJ_API_KEY)
                .modelName(ApiKey.GJ_EMBEDDING_MODEL)
                .build();


        //这里就是将文本转为向量
        Response<Embedding> embeddingResponse = embeddingModel.embed("你好 我是幽兰 ");
        System.out.println(embeddingResponse.content());
        System.out.println(embeddingResponse.content().vector().length);

        //将向量存储到向量数据库中
        /*ClickHouseSettings clickHouseSettings = ClickHouseSettings.builder()
                .url("http://47.109.73.232:8123")
                .database("default")
                .username("default")
                .password("LonglongTz1")
                .dimension(1024)
                .build();

        EmbeddingStore embeddingStore = ClickHouseEmbeddingStore.builder()
                .settings(clickHouseSettings)
                .build();

        String add = embeddingStore.add(embeddingResponse.content());
        System.out.println(add);
*/

    /*    //使用 zilliz 向量数据库
        EmbeddingStore embeddingStore = MilvusEmbeddingStore.builder()
                .host("in01-70987aa9af97097.ali-cn-hangzhou.vectordb.zilliz.com.cn")
                .port(19530)
                .username("db_admin")
                .password("Tn6|Zf;ds?J/tGKu")
                .collectionName("test")
                .dimension(1024)
                .build();
        embeddingStore.add(embeddingResponse.content());*/

        EmbeddingStore<TextSegment> embeddingStore = PgVectorEmbeddingStore.builder()
                .host("localhost")
                .port(5433)
                .database("postgres")
                .table("xp_test_vector")
                .user("root")
                .password("123456")
                .dimension(1024)
                .build();

        //这里是存储向量
//        embeddingStore.add(embeddingResponse.content());

        //对一个文本向量化
        Response<Embedding> embeddingR1 = embeddingModel.embed("今天天气是真的不错");

        //接下来就是对向量的相似度匹配
        EmbeddingSearchRequest embeddingSearchRequest = EmbeddingSearchRequest.builder()
                .queryEmbedding(embeddingR1.content())
                .maxResults(10)
                .minScore(0.0)
                .build();
        EmbeddingSearchResult<TextSegment> searchResult = embeddingStore.search(embeddingSearchRequest);
        List<EmbeddingMatch<TextSegment>> matches = searchResult.matches();
        for (EmbeddingMatch<TextSegment> match : matches) {
            System.out.println("match.score() = " + match.score());
        }

        System.out.println("存储成功");

    }
}
